﻿using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Diagnostics;
using System.IO;
using System.Runtime.InteropServices;
using System.Text;

namespace SmartBot
{
    /// <summary>
    /// Class for DLL workers
    /// </summary>
    public class DLLInjectHelper
    {
        /// <summary>
        /// Result for DLL injection
        /// </summary>
        public enum DLLInjectResults
        {
            DllNotFound = 0, 
            ProcessNotFound = 1, 
            Failed = 2, 
            Success = 3, 
            AccessIsDenied = 4, 
            NoLoadLibraryA = 5, 
            CannotAllocateMemory = 6, 
            CannotWriteDLLToMemory = 7, 
            VersionMismatch = 8, 
            Others = 9
        }

        /// <summary>
        /// Equal to NULL in C++
        /// </summary>
        public static readonly IntPtr INTPTR_ZERO = (IntPtr)0;

        #region Imported functions from DLL
        [DllImport("kernel32.dll", SetLastError = true)]
        public static extern IntPtr OpenProcess(uint dwDesiredAccess, int bInheritHandle, uint dwProcessId);

        [DllImport("kernel32.dll", SetLastError = true, CallingConvention = CallingConvention.Winapi)]
        [return: MarshalAs(UnmanagedType.Bool)]
        public static extern bool IsWow64Process([In] IntPtr processHandle, [Out, MarshalAs(UnmanagedType.Bool)] out bool wow64Process);

        [DllImport("kernel32.dll", SetLastError = true)]
        public static extern int CloseHandle(IntPtr hObject);

        [DllImport("kernel32.dll", SetLastError = true)]
        public static extern uint GetLastError();

        [DllImport("kernel32.dll", SetLastError = true)]
        public static extern IntPtr GetProcAddress(IntPtr hModule, string lpProcName);

        [DllImport("kernel32.dll", SetLastError = true)]
        public static extern IntPtr GetModuleHandle(string lpModuleName);

        [DllImport("kernel32.dll", SetLastError = true)]
        public static extern IntPtr VirtualAllocEx(IntPtr hProcess, IntPtr lpAddress, IntPtr dwSize, uint flAllocationType, uint flProtect);

        [DllImport("kernel32.dll", SetLastError = true)]
        public static extern int WriteProcessMemory(IntPtr hProcess, IntPtr lpBaseAddress, byte[] buffer, uint size, int lpNumberOfBytesWritten);

        [DllImport("kernel32.dll", SetLastError = true)]
        public static extern IntPtr CreateRemoteThread(IntPtr hProcess, IntPtr lpThreadAttribute, IntPtr dwStackSize, IntPtr lpStartAddress,
            IntPtr lpParameter, uint dwCreationFlags, IntPtr lpThreadId);

        //Support for Shared Name Pipe
        [DllImport("kernel32.dll", SetLastError = true, CharSet = CharSet.Auto)]
        public static extern IntPtr OpenFileMapping(uint dwDesiredAccess, bool bInheritHandle, string lpName);

        [DllImport("msvcrt.dll", EntryPoint = "memcpy", CallingConvention = CallingConvention.Cdecl, SetLastError = false)]
        public static extern IntPtr MemCopy(IntPtr dest, IntPtr src, uint count);

        public const UInt32 STANDARD_RIGHTS_REQUIRED = 0x000F0000;
        public const UInt32 SECTION_QUERY = 0x0001;
        public const UInt32 SECTION_MAP_WRITE = 0x0002;
        public const UInt32 SECTION_MAP_READ = 0x0004;
        public const UInt32 SECTION_MAP_EXECUTE = 0x0008;
        public const UInt32 SECTION_EXTEND_SIZE = 0x0010;
        public const UInt32 SECTION_ALL_ACCESS = (STANDARD_RIGHTS_REQUIRED | SECTION_QUERY |
            SECTION_MAP_WRITE |
            SECTION_MAP_READ |
            SECTION_MAP_EXECUTE |
            SECTION_EXTEND_SIZE);
        public const UInt32 FILE_MAP_ALL_ACCESS = SECTION_ALL_ACCESS;

        [DllImport("kernel32.dll", EntryPoint = "CopyMemory", SetLastError = false)]
        public static extern void CopyMemory(IntPtr dest, IntPtr src, uint count);

        [DllImport("kernel32.dll", SetLastError = true)]
        public static extern IntPtr MapViewOfFile(
            IntPtr hFileMappingObject,
            //FileMapAccess dwDesiredAccess,
            UInt32 dwDesiredAccess,
            UInt32 dwFileOffsetHigh,
            UInt32 dwFileOffsetLow,
            UIntPtr dwNumberOfBytesToMap);

        #endregion

        /// <summary>
        /// Inject the DLL to remote process, specifying the DLL path.
        /// </summary>
        /// <param name="processName">Process name as string</param>
        /// <param name="dllPath">DLL path</param>
        /// <returns></returns>
        public DLLInjectResults InjectDLL(string processName, string dllPath)
        {
            //Determine if the DLL file exists, otherwise, return not found error
            if (!File.Exists(dllPath))
            {
                return DLLInjectResults.DllNotFound;
            }

            uint _processID = 0;

            //List all processes on the computer
            Process[] _allProcesses = Process.GetProcesses();

            bool is64Bit = false;
            foreach (Process proc in _allProcesses)
            {
                //if (proc.MainWindowTitle == "Untitled - Paint")
                if (proc.ProcessName == processName)
                {
                    _processID = (uint)proc.Id;
                    is64Bit = IsRemoteProcess64Bit(proc);

                    break;
                }
            }

            if (_processID == 0)
            {
                return DLLInjectResults.ProcessNotFound;
            }

            //Call the real injection process
            return ExecuteDLLInjection(_processID, dllPath, is64Bit);
        }

        /// <summary>
        /// Check if the remote process is a 64 bit.
        /// </summary>
        /// <param name="myProcess"></param>
        /// <returns></returns>
        public bool IsRemoteProcess64Bit(Process myProcess)
        {
            /*//No need to check if the OS is 64 bit
            if (!Environment.Is64BitOperatingSystem)
                return false;

            bool is64bit = false;
            if (!IsWow64Process(myProcess.Handle, out is64bit))
                throw new Win32Exception(Marshal.GetLastWin32Error());

            return !is64bit;
             */
            return false;
        }

        /// <summary>
        /// Detect if the current process is 64 bit or 32 bit
        /// </summary>
        /// <returns></returns>
        public bool IsMyself64Bit()
        {
            //return Environment.Is64BitProcess;
            return false;
        }

        /// <summary>
        /// Do the real injection to the remote process
        /// </summary>
        /// <param name="processID">Process ID to inject</param>
        /// <param name="dllPath">Path to the DLL file</param>
        /// <returns>Status message</returns>
        DLLInjectResults ExecuteDLLInjection(uint processID, string dllPath, bool use64Bit)
        {

            //Will fail if proceed.
            if (use64Bit == true && IsMyself64Bit() == false)
            {
                return DLLInjectResults.VersionMismatch;
            }

            IntPtr hndProc = OpenProcess((0x2 | 0x8 | 0x10 | 0x20 | 0x400), 1, processID);

            if (hndProc == INTPTR_ZERO)
            {
                return DLLInjectResults.ProcessNotFound;
            }

            //Định vị LoadLibraryA function - library này được load tại cùng một vị trí memory trên tất cả các processes.
            //Giá trị trả về sẽ luôn đúng khi áp dụng trên toàn bộ các process khác nhau.
            IntPtr lpLLAddress = GetProcAddress(GetModuleHandle("kernel32.dll"), "LoadLibraryA");

            //Check if the reference address is not found
            if (lpLLAddress == INTPTR_ZERO)
            {
                return DLLInjectResults.NoLoadLibraryA;
            }

            //Assign memory of the remote process
            IntPtr lpAddress = VirtualAllocEx(hndProc, (IntPtr)null, (IntPtr)dllPath.Length, (0x1000 | 0x2000), 0X40);

            if (lpAddress == INTPTR_ZERO)
            {
                return DLLInjectResults.CannotAllocateMemory;
            }

            byte[] bytes = Encoding.ASCII.GetBytes(dllPath);

            //Write the path of the DLL to remote process so it can call the DLL remotely
            if (WriteProcessMemory(hndProc, lpAddress, bytes, (uint)bytes.Length, 0) == 0)
            {
                return DLLInjectResults.CannotWriteDLLToMemory;
            }

            //Sử dụng CreateRemoteThread để inject DLL vào remote process
            //Phân biệt remote process là 64 bit hay 32 bit
            if (CreateRemoteThread(hndProc, (IntPtr)null, INTPTR_ZERO, lpLLAddress, lpAddress, 0, (IntPtr)null) == INTPTR_ZERO)
            {
                int errorID = new Win32Exception(Marshal.GetLastWin32Error()).NativeErrorCode;
                string message = new Win32Exception(Marshal.GetLastWin32Error()).Message;

                //Write error to log file
                if (errorID == 5)
                    return DLLInjectResults.AccessIsDenied;
            }

            //Try to free the DLL right away

            CloseHandle(hndProc);

            return DLLInjectResults.Success;
        }
    }
}
